#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#ifdef _WIN32
#include <winsock2.h>
#include <iphlpapi.h>
#include <icmpapi.h>
#include <signal.h>
#pragma comment(lib, "iphlpapi.lib")
#pragma comment(lib, "ws2_32.lib")
#else
#include <sys/time.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netdb.h>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <unistd.h>
#include <signal.h>
#include <getopt.h>  // 仅在非MSVC平台包含
#endif

#define THREAD_POOL_IMPLEMENTATION
#include "tpool.hol.h"
#define ASYNC_IMPLEMENTATION
#include "async.hol.h"


// 默认配置
#define DEFAULT_PACKET_SIZE 64
#define DEFAULT_PING_COUNT 4
#define DEFAULT_TIMEOUT 1000  // 毫秒
#define MAX_WAIT_AFTER_INTERRUPT 3000  // 中断后最大等待时间(毫秒)

// 全局变量用于处理程序中断
static volatile bool g_should_stop = false;
static volatile time_t g_interrupt_time = 0;  // 记录中断发生的时间

// 清理URL中的协议前缀（http://、https://）
const char* clean_hostname(const char* input) {
    static char cleaned[256];
    strncpy(cleaned, input, sizeof(cleaned)-1);
    
    // 去除 http:// 前缀
    if (strstr(cleaned, "http://") == cleaned) {
        memmove(cleaned, cleaned + 7, strlen(cleaned) - 6);
    }
    // 去除 https:// 前缀
    else if (strstr(cleaned, "https://") == cleaned) {
        memmove(cleaned, cleaned + 8, strlen(cleaned) - 7);
    }
    
    return cleaned;
}

// Ping结果结构体
typedef struct {
    char host[256];          // 主机名（原始输入）
    char clean_host[256];    // 清理后的主机名（用于解析）
    char ip[16];             // IP地址
    int seq;                 // 序列号
    float rtt;               // 往返时间(毫秒)
    bool success;            // 是否成功
} ping_result_t;

// 配置参数结构体
typedef struct {
    int count;               // 发送次数
    int packet_size;         // 数据包大小
    int timeout;             // 超时时间(毫秒)
    int num_hosts;           // 主机数量
    char** hosts;            // 主机列表（原始输入）
    char** clean_hosts;      // 清理后的主机列表（用于解析）
} ping_config_t;

// 计算校验和
unsigned short calculate_checksum(void* data, int len) {
    unsigned short* buf = (unsigned short*)data;
    unsigned int sum = 0;
    unsigned short result;

    for (sum = 0; len > 1; len -= 2)
        sum += *buf++;
    if (len == 1)
        sum += *(unsigned char*)buf;
    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    result = ~sum;
    return result;
}

// 信号处理函数，用于捕获Ctrl+C
void handle_sigint(int sig) {
    if (!g_should_stop) {  // 只处理第一次中断
        g_should_stop = true;
        g_interrupt_time = time(NULL);
        printf("\n收到中断信号，正在整理结果...\n");
        printf("（如果长时间无响应，请再次按下Ctrl+C强制退出）\n");
    } else {  // 第二次中断，强制退出
        printf("\n收到第二次中断信号，强制退出...\n");
        exit(1);
    }
}

// Windows平台控制台事件处理函数
#ifdef _WIN32
BOOL WINAPI ConsoleHandler(DWORD dwCtrlType) {
    if (dwCtrlType == CTRL_C_EVENT) {
        handle_sigint(SIGINT);
        return TRUE;  // 表示已处理事件
    }
    return FALSE;
}
#endif

// 检查是否应该退出（考虑超时）
bool should_exit() {
    if (!g_should_stop) return false;
    
    // 如果超过最大等待时间，强制退出
    if (time(NULL) - g_interrupt_time > MAX_WAIT_AFTER_INTERRUPT / 1000) {
        printf("\n等待超时，强制退出...\n");
        exit(1);
    }
    return true;
}

// Ping任务函数
void* ping_host(void* args) {
    // 检查是否需要立即终止任务
    if (g_should_stop) {
        free(args);
        return NULL;
    }
    
    // 参数是一个包含主机名、序列号和配置的结构体
    struct {
        char host[256];        // 原始主机名（用于显示）
        char clean_host[256];  // 清理后的主机名（用于解析）
        int seq;
        ping_config_t* config;
    }* param = (void*)args;
    
    ping_result_t* result = malloc(sizeof(ping_result_t));
    memset(result, 0, sizeof(ping_result_t));
    strncpy(result->host, param->host, sizeof(result->host) - 1);
    strncpy(result->clean_host, param->clean_host, sizeof(result->clean_host) - 1);
    result->seq = param->seq;
    result->success = false;

#ifdef _WIN32
    // Windows实现
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        free(param);
        return result;
    }

    HANDLE hIcmpFile = IcmpCreateFile();
    if (hIcmpFile == INVALID_HANDLE_VALUE) {
        WSACleanup();
        free(param);
        return result;
    }

    // 使用清理后的主机名解析
    struct hostent* he = gethostbyname(param->clean_host);
    if (!he) {
        IcmpCloseHandle(hIcmpFile);
        WSACleanup();
        free(param);
        return result;
    }

    strcpy(result->ip, inet_ntoa(*(struct in_addr*)he->h_addr_list[0]));
    
    char* SendData = malloc(param->config->packet_size);
    memset(SendData, 'P', param->config->packet_size); // 填充数据包
    
    LPVOID ReplyBuffer = malloc(sizeof(ICMP_ECHO_REPLY) + param->config->packet_size);
    
    // 检查是否需要在发送前终止
    if (should_exit()) {
        free(SendData);
        free(ReplyBuffer);
        IcmpCloseHandle(hIcmpFile);
        WSACleanup();
        free(param);
        free(result);
        return NULL;
    }
    
    DWORD dwRetVal = IcmpSendEcho(hIcmpFile, 
                                 inet_addr(result->ip),
                                 SendData, param->config->packet_size,
                                 NULL, ReplyBuffer, 
                                 sizeof(ICMP_ECHO_REPLY) + param->config->packet_size, 
                                 param->config->timeout);

    // 检查是否在等待响应时收到中断
    if (should_exit()) {
        free(SendData);
        free(ReplyBuffer);
        IcmpCloseHandle(hIcmpFile);
        WSACleanup();
        free(param);
        free(result);
        return NULL;
    }

    if (dwRetVal > 0) {
        PICMP_ECHO_REPLY pEchoReply = (PICMP_ECHO_REPLY)ReplyBuffer;
        result->success = true;
        result->rtt = pEchoReply->RoundTripTime;
    }

    free(SendData);
    free(ReplyBuffer);
    IcmpCloseHandle(hIcmpFile);
    WSACleanup();
#else
    // Linux实现
    int sockfd;
    struct sockaddr_in dest_addr;
    struct hostent* he;
    char* packet = malloc(param->config->packet_size);
    struct icmphdr* icmp_hdr;
    struct timeval start, end;

    // 检查是否需要立即终止
    if (should_exit()) {
        free(packet);
        free(param);
        free(result);
        return NULL;
    }

    // 创建原始套接字
    if ((sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP)) < 0) {
        free(packet);
        free(param);
        return result;
    }

    // 使用清理后的主机名解析
    if ((he = gethostbyname(param->clean_host)) == NULL) {
        close(sockfd);
        free(packet);
        free(param);
        return result;
    }

    // 设置目标地址
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    memcpy(&dest_addr.sin_addr.s_addr, he->h_addr, he->h_length);
    strcpy(result->ip, inet_ntoa(dest_addr.sin_addr));

    // 构建ICMP包
    memset(packet, 0, param->config->packet_size);
    icmp_hdr = (struct icmphdr*)packet;
    icmp_hdr->type = ICMP_ECHO;
    icmp_hdr->code = 0;
    icmp_hdr->un.echo.id = getpid() & 0xFFFF;
    icmp_hdr->un.echo.sequence = param->seq;
    
    // 填充数据包内容
    memset(packet + sizeof(struct icmphdr), 'P', 
           param->config->packet_size - sizeof(struct icmphdr));
    
    icmp_hdr->checksum = calculate_checksum(packet, param->config->packet_size);

    // 检查是否需要在发送前终止
    if (should_exit()) {
        close(sockfd);
        free(packet);
        free(param);
        free(result);
        return NULL;
    }

    // 发送ICMP请求
    gettimeofday(&start, NULL);
    if (sendto(sockfd, packet, param->config->packet_size, 0, 
              (struct sockaddr*)&dest_addr, sizeof(dest_addr)) <= 0) {
        close(sockfd);
        free(packet);
        free(param);
        return result;
    }

    // 接收响应
    struct sockaddr_in src_addr;
    socklen_t addr_len = sizeof(src_addr);
    char recv_buf[1024];
    
    fd_set readfds;
    struct timeval timeout;
    timeout.tv_sec = param->config->timeout / 1000;
    timeout.tv_usec = (param->config->timeout % 1000) * 1000;
    
    FD_ZERO(&readfds);
    FD_SET(sockfd, &readfds);

    // 等待响应时定期检查是否需要退出
    int ret;
    while (true) {
        // 设置一个短超时，以便定期检查中断信号
        struct timeval short_timeout = {0, 100000}; // 100ms
        ret = select(sockfd + 1, &readfds, NULL, NULL, &short_timeout);
        
        if (should_exit()) { // 检查是否需要退出
            close(sockfd);
            free(packet);
            free(param);
            free(result);
            return NULL;
        }
        
        if (ret > 0) break; // 有数据可读
        if (--timeout.tv_sec < 0 && timeout.tv_usec <= 0) break; // 超时
    }

    if (ret > 0 && FD_ISSET(sockfd, &readfds)) {
        ssize_t len = recvfrom(sockfd, recv_buf, sizeof(recv_buf), 0,
                              (struct sockaddr*)&src_addr, &addr_len);
        if (len > 0) {
            // 验证是否是ICMP响应
            struct ip* ip_hdr = (struct ip*)recv_buf;
            int ip_header_len = ip_hdr->ip_hl * 4;
            struct icmphdr* recv_icmp_hdr = (struct icmphdr*)(recv_buf + ip_header_len);
            
            // 检查是否是我们发送的请求的响应
            if (recv_icmp_hdr->type == ICMP_ECHOREPLY &&
                recv_icmp_hdr->un.echo.id == (getpid() & 0xFFFF) &&
                recv_icmp_hdr->un.echo.sequence == param->seq) {
                
                gettimeofday(&end, NULL);
                // 计算往返时间(毫秒)
                result->rtt = (end.tv_sec - start.tv_sec) * 1000.0 +
                             (end.tv_usec - start.tv_usec) / 1000.0;
                result->success = true;
            }
        }
    }

    close(sockfd);
    free(packet);
#endif

    free(param);
    return result;
}

// 打印使用帮助
void print_usage(const char* program_name) {
    printf("用法: %s [选项] <主机1> [主机2] ...\n", program_name);
    printf("跨平台多主机ping工具\n\n");
    printf("选项:\n");
    printf("  -c <次数>    设置每个主机发送的ping包数量 (默认: %d)\n", DEFAULT_PING_COUNT);
    printf("  -s <大小>    设置ping包大小(字节) (默认: %d)\n", DEFAULT_PACKET_SIZE);
    printf("  -t <超时>    设置超时时间(毫秒) (默认: %d)\n", DEFAULT_TIMEOUT);
    printf("  -h           显示帮助信息\n");
}

// MSVC专用的命令行参数解析（替代getopt）
#ifdef _WIN32
bool parse_args_msvc(int argc, char* argv[], ping_config_t* config) {
    int num_hosts = 0;
    char** hosts = NULL;

    // 设置默认值
    config->count = DEFAULT_PING_COUNT;
    config->packet_size = DEFAULT_PACKET_SIZE;
    config->timeout = DEFAULT_TIMEOUT;

    // 解析参数
    for (int i = 1; i < argc; i++) {
        // 处理选项参数
        if (argv[i][0] == '-') {
            switch (argv[i][1]) {
                case 'c':  // -c <次数>
                    if (i + 1 < argc) {
                        config->count = atoi(argv[++i]);
                        if (config->count <= 0) {
                            fprintf(stderr, "错误: 次数必须是正数\n");
                            return false;
                        }
                    } else {
                        fprintf(stderr, "错误: -c 需要指定次数\n");
                        return false;
                    }
                    break;
                    
                case 's':  // -s <大小>
                    if (i + 1 < argc) {
                        config->packet_size = atoi(argv[++i]);
                        if (config->packet_size < 1 || config->packet_size > 65507) {
                            fprintf(stderr, "错误: 包大小必须在1-65507之间\n");
                            return false;
                        }
                    } else {
                        fprintf(stderr, "错误: -s 需要指定大小\n");
                        return false;
                    }
                    break;
                    
                case 't':  // -t <超时>
                    if (i + 1 < argc) {
                        config->timeout = atoi(argv[++i]);
                        if (config->timeout <= 0) {
                            fprintf(stderr, "错误: 超时时间必须是正数\n");
                            return false;
                        }
                    } else {
                        fprintf(stderr, "错误: -t 需要指定超时时间\n");
                        return false;
                    }
                    break;
                    
                case 'h':  // -h 帮助
                    print_usage(argv[0]);
                    return false;
                    
                default:  // 未知选项
                    fprintf(stderr, "错误: 未知选项 %s\n", argv[i]);
                    fprintf(stderr, "使用 -h 查看帮助\n");
                    return false;
            }
        } else {
            // 收集主机名
            hosts = realloc(hosts, sizeof(char*) * (num_hosts + 1));
            hosts[num_hosts++] = argv[i];
        }
    }

    // 检查是否有主机名
    if (num_hosts <= 0) {
        fprintf(stderr, "错误: 至少需要指定一个主机\n");
        fprintf(stderr, "使用 -h 查看帮助\n");
        free(hosts);
        return false;
    }

    // 保存主机名并清理前缀
    config->num_hosts = num_hosts;
    config->hosts = hosts;
    config->clean_hosts = malloc(sizeof(char*) * num_hosts);
    for (int i = 0; i < num_hosts; i++) {
        config->clean_hosts[i] = strdup(clean_hostname(hosts[i]));
    }

    return true;
}
#endif

// 非MSVC平台使用getopt解析参数
#ifndef _WIN32
bool parse_args(int argc, char* argv[], ping_config_t* config) {
    int opt;
    
    // 设置默认值
    config->count = DEFAULT_PING_COUNT;
    config->packet_size = DEFAULT_PACKET_SIZE;
    config->timeout = DEFAULT_TIMEOUT;
    config->num_hosts = 0;
    config->hosts = NULL;
    config->clean_hosts = NULL;  // 初始化清理后的主机列表
    
    // 解析选项
    while ((opt = getopt(argc, argv, "c:s:t:h")) != -1) {
        switch (opt) {
            case 'c':
                config->count = atoi(optarg);
                if (config->count <= 0) {
                    fprintf(stderr, "错误: 次数必须是正数\n");
                    return false;
                }
                break;
            case 's':
                config->packet_size = atoi(optarg);
                if (config->packet_size < 1 || config->packet_size > 65507) {
                    fprintf(stderr, "错误: 包大小必须在1-65507之间\n");
                    return false;
                }
                break;
            case 't':
                config->timeout = atoi(optarg);
                if (config->timeout <= 0) {
                    fprintf(stderr, "错误: 超时时间必须是正数\n");
                    return false;
                }
                break;
            case 'h':
                print_usage(argv[0]);
                return false;
            default:
                fprintf(stderr, "使用 -h 查看帮助\n");
                return false;
        }
    }
    
    // 收集主机名并清理
    config->num_hosts = argc - optind;
    if (config->num_hosts <= 0) {
        fprintf(stderr, "错误: 至少需要指定一个主机\n");
        fprintf(stderr, "使用 -h 查看帮助\n");
        return false;
    }
    
    // 保存原始主机名和清理后的主机名
    config->hosts = malloc(sizeof(char*) * config->num_hosts);
    config->clean_hosts = malloc(sizeof(char*) * config->num_hosts);
    for (int i = 0; i < config->num_hosts; i++) {
        config->hosts[i] = argv[optind + i];
        config->clean_hosts[i] = strdup(clean_hostname(argv[optind + i]));  // 清理前缀
    }
    
    return true;
}
#endif

// 并行Ping多个主机
void ping_hosts(ping_config_t* config) {
    // 注册信号处理函数（跨平台实现）
#ifdef _WIN32
    if (!SetConsoleCtrlHandler(ConsoleHandler, TRUE)) {
        fprintf(stderr, "设置控制台事件处理失败\n");
    }
#else
    struct sigaction sa;
    sa.sa_handler = handle_sigint;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = 0;
    sigaction(SIGINT, &sa, NULL);
#endif

    async_init();
    
    // 计算总任务数
    int total_tasks = config->num_hosts * config->count;
    async_result_t** results = malloc(sizeof(async_result_t*) * total_tasks);
    if (!results) {
        fprintf(stderr, "内存分配失败\n");
        return;
    }
    
    // 显示开始信息
    printf("正在ping %d个主机，每个主机发送%d个数据包...\n", 
           config->num_hosts, config->count);
    printf("数据包大小: %d字节，超时时间: %d毫秒\n\n",
           config->packet_size, config->timeout);
    
    // 为每个主机创建ping任务（使用清理后的主机名）
    int task_idx = 0;
    for (int i = 0; i < config->num_hosts && !g_should_stop; i++) {
        for (int seq = 1; seq <= config->count && !g_should_stop; seq++) {
            // 创建任务参数
            struct {
                char host[256];        // 原始主机名（显示用）
                char clean_host[256];  // 清理后的主机名（解析用）
                int seq;
                ping_config_t* config;
            }* param = malloc(sizeof(*param));
            
            if (!param) {
                fprintf(stderr, "内存分配失败\n");
                g_should_stop = true;
                break;
            }
            
            strncpy(param->host, config->hosts[i], sizeof(param->host) - 1);
            strncpy(param->clean_host, config->clean_hosts[i], sizeof(param->clean_host) - 1);
            param->seq = seq;
            param->config = config;
            
            // 提交异步任务
            results[task_idx++] = async(ping_host, param);
        }
    }
    
    // 收集并统计结果
    struct host_stats {
        int sent;
        int received;
        float min_rtt;
        float max_rtt;
        float avg_rtt;
        float total_rtt;
    }* stats = calloc(config->num_hosts, sizeof(struct host_stats));
    
    if (!stats) {
        fprintf(stderr, "内存分配失败\n");
        return;
    }
    
    // 初始化统计信息
    for (int i = 0; i < config->num_hosts; i++) {
        stats[i].min_rtt = 1e9; // 一个很大的数
    }
    
    // 处理结果，增加中断检查
    for (int i = 0; i < task_idx; i++) {
        // 检查是否需要退出
        if (should_exit()) break;
        
        // 非阻塞等待结果，每100ms检查一次是否需要退出
        ping_result_t* res = NULL;
        while (!res && !should_exit()) {
            res = (ping_result_t*)async_try_await(results[i]);
            if (!res) {
#ifdef _WIN32
                Sleep(100);
#else
                usleep(100000); // 100ms
#endif
            }
        }
        
        if (res) {
            // 找到对应的主机索引
            int host_idx = -1;
            for (int j = 0; j < config->num_hosts; j++) {
                if (strcmp(res->host, config->hosts[j]) == 0) {
                    host_idx = j;
                    break;
                }
            }
            
            if (host_idx != -1) {
                stats[host_idx].sent++;
                
                // 显示结果（保留原始输入，如带https的URL）
                if (res->success) {
                    stats[host_idx].received++;
                    stats[host_idx].total_rtt += res->rtt;
                    stats[host_idx].avg_rtt = stats[host_idx].total_rtt / stats[host_idx].received;
                    
                    if (res->rtt < stats[host_idx].min_rtt) {
                        stats[host_idx].min_rtt = res->rtt;
                    }
                    if (res->rtt > stats[host_idx].max_rtt) {
                        stats[host_idx].max_rtt = res->rtt;
                    }
                    
                    printf("%s (%s): 回复 from %s: 字节=%d 时间=%.2fms TTL=? 序列号=%d\n",
                           res->host, res->ip, res->ip,
                           config->packet_size, res->rtt, res->seq);
                } else {
                    printf("%s: 请求超时。序列号=%d\n", res->host, res->seq);
                }
            }
            
            free(res);
        }
    }
    
    // 打印每个主机的统计信息
    printf("\n--- 统计信息 ---\n");
    for (int i = 0; i < config->num_hosts; i++) {
        if (stats[i].sent == 0) continue;
        
        float loss_rate = (stats[i].sent - stats[i].received) * 100.0 / stats[i].sent;
        
        printf("\n%s 的ping统计信息:\n", config->hosts[i]);
        printf("    数据包: 已发送 = %d, 已接收 = %d, 丢失 = %d (%.1f%% 丢失)\n",
               stats[i].sent, stats[i].received,
               stats[i].sent - stats[i].received, loss_rate);
        
        if (stats[i].received > 0) {
            printf("    往返时间: 最小 = %.2fms, 最大 = %.2fms, 平均 = %.2fms\n",
                   stats[i].min_rtt, stats[i].max_rtt, stats[i].avg_rtt);
        }
    }
    
    // 清理资源
    free(stats);
    free(results);
    async_cleanup();
}

int main(int argc, char* argv[]) {
    ping_config_t config;
    bool parse_success = false;
    
    // 根据平台选择参数解析方式
#ifdef _WIN32
    parse_success = parse_args_msvc(argc, argv, &config);
#else
    parse_success = parse_args(argc, argv, &config);
#endif
    
    // 解析失败处理
    if (!parse_success) {
        if (config.hosts) free(config.hosts);
        if (config.clean_hosts) {
            for (int i = 0; i < config.num_hosts; i++) {
                free(config.clean_hosts[i]);
            }
            free(config.clean_hosts);
        }
        return 1;
    }
    
    // 执行ping操作
    ping_hosts(&config);
    
    // 清理资源
    free(config.hosts);
    for (int i = 0; i < config.num_hosts; i++) {
        free(config.clean_hosts[i]);
    }
    free(config.clean_hosts);
    return 0;
}